#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
主窗口类 - PyQt6版本
对应原C++项目的MainWindow类
"""

import sys
import os
import cv2
import numpy as np
import threading
import time
import yaml
from datetime import datetime

from PySide6.QtWidgets import QMainWindow, QLabel, QPushButton, QFileDialog, QVBoxLayout, QHBoxLayout, QWidget, QMessageBox
from PySide6.QtCore import Qt, QTimer, Signal, Slot, QThread
from PySide6.QtGui import QPixmap, QImage, QPainter, QColor, QFont

# 导入自定义模块
from xdma_driver import XDMADriver
from image_processor import YOLOv5Detector
from threads import SendImage, EventThread, EventThread1, EventThread2, LoadImageThread, StartCalcThread


class MainWindow(QMainWindow):
    """主窗口类"""

    def __init__(self):
        super().__init__()

        # 初始化成员变量
        self.img_height = 720
        self.img_width = 1280
        self.cam_addr_0 = 0x30000000
        self.cam_addr_1 = 0x31000000

        # YOLOv5相关参数
        self.image_h = 640
        self.image_w = 640
        self.class_num = 20
        self.class_names = [
            "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat",
            "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person",
            "pottedplant", "sheep", "sofa", "train", "tvmonitor"
        ]
        self.channel = (self.class_num + 5) * 3

        # NMS参数
        self.nms_zero = [84, 89, 99]
        self.nms_scale = [0.11205127090215683, 0.11795272678136826, 0.12487365305423737]

        # 网格尺寸
        self.h_0, self.w_0 = 20, 20
        self.h_1, self.w_1 = 40, 40
        self.h_2, self.w_2 = 80, 80
        self.nms_total = (self.h_0 * self.w_0 + self.h_1 * self.w_1 + self.h_2 * self.w_2) * 3

        # 置信度和IOU阈值
        self.conf_thres = 0.3
        self.iou_thres = 0.5

        # Anchor boxes
        self.anchor_w = [10, 16, 33, 30, 62, 59, 116, 156, 373]
        self.anchor_h = [13, 30, 23, 61, 45, 119, 90, 198, 326]

        # 计算步长和anchor
        self._calculate_strides_and_anchors()

        # 初始化组件
        self.xdma_driver = None
        self.yolov5_detector = None

        # 初始化线程和图像处理
        self.send_image = SendImage()
        self.event_thread = None
        self.event_thread1 = None
        self.event_thread2 = None
        self.load_img_thread = None
        self.start_calc_thread = None

        # 标志位
        self.collect_flag = False
        self.start_flag = False

        # 图像显示相关
        self.current_image = None

        # 初始化UI
        self._init_ui()

        # 初始化硬件连接
        self._init_hardware()

    def _calculate_strides_and_anchors(self):
        """计算步长和anchor boxes"""
        # 计算步长
        self.stride_h_0 = self.image_h / self.h_0
        self.stride_w_0 = self.image_w / self.w_0
        self.stride_h_1 = self.image_h / self.h_1
        self.stride_w_1 = self.image_w / self.w_1
        self.stride_h_2 = self.image_h / self.h_2
        self.stride_w_2 = self.image_w / self.w_2

        # 计算anchor boxes (相对于特征图的比例)
        self.anchor_w_0 = [
            self.anchor_w[6] / self.stride_w_0,
            self.anchor_w[7] / self.stride_w_0,
            self.anchor_w[8] / self.stride_w_0
        ]
        self.anchor_h_0 = [
            self.anchor_h[6] / self.stride_h_0,
            self.anchor_h[7] / self.stride_h_0,
            self.anchor_h[8] / self.stride_h_0
        ]

        self.anchor_w_1 = [
            self.anchor_w[3] / self.stride_w_1,
            self.anchor_w[4] / self.stride_w_1,
            self.anchor_w[5] / self.stride_w_1
        ]
        self.anchor_h_1 = [
            self.anchor_h[3] / self.stride_h_1,
            self.anchor_h[4] / self.stride_h_1,
            self.anchor_h[5] / self.stride_h_1
        ]

        self.anchor_w_2 = [
            self.anchor_w[0] / self.stride_w_2,
            self.anchor_w[1] / self.stride_w_2,
            self.anchor_w[2] / self.stride_w_2
        ]
        self.anchor_h_2 = [
            self.anchor_h[0] / self.stride_h_2,
            self.anchor_h[1] / self.stride_h_2,
            self.anchor_h[2] / self.stride_h_2
        ]

    def _init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("FPGA开源工坊——YOLOv5相机检测")
        self.setFixedSize(self.img_width + 200, self.img_height + 200)

        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 创建布局
        layout = QVBoxLayout(central_widget)

        # 图像显示区域
        self.image_label = QLabel()
        self.image_label.setFixedSize(self.img_width, self.img_height)
        self.image_label.setStyleSheet("QLabel { background-color: #000000; }")
        layout.addWidget(self.image_label)

        # 按钮布局
        button_layout = QHBoxLayout()

        # 启动摄像头按钮
        self.start_camera_btn = QPushButton("启动摄像头")
        self.start_camera_btn.clicked.connect(self.on_start_camera_clicked)
        button_layout.addWidget(self.start_camera_btn)

        # 加载指令按钮
        self.load_instruction_btn = QPushButton("加载指令")
        self.load_instruction_btn.clicked.connect(self.on_load_instruction_clicked)
        button_layout.addWidget(self.load_instruction_btn)

        # 加载权重按钮
        self.load_weight_btn = QPushButton("加载权重")
        self.load_weight_btn.clicked.connect(self.on_load_weight_clicked)
        button_layout.addWidget(self.load_weight_btn)

        layout.addLayout(button_layout)

        # 状态信息显示
        self.status_label = QLabel("就绪")
        self.status_label.setStyleSheet("QLabel { color: blue; font-size: 12px; }")
        layout.addWidget(self.status_label)

    def _init_hardware(self):
        """初始化硬件连接"""
        try:
            # 初始化XDMA驱动
            self.xdma_driver = XDMADriver()

            # 初始化YOLOv5检测器
            self.yolov5_detector = YOLOv5Detector()

            # 创建事件线程
            self.event_thread = EventThread(self.xdma_driver)
            self.event_thread1 = EventThread1(self.xdma_driver)
            self.event_thread2 = EventThread2(self.xdma_driver)

            # 创建图像处理线程
            self.load_img_thread = LoadImageThread(self.send_image)
            self.start_calc_thread = StartCalcThread(self.xdma_driver, self.send_image)

            # 连接信号和槽
            self.event_thread.sig_get_nms.connect(self.on_nms)
            self.event_thread1.sig_get_img.connect(self.on_get_img)
            self.event_thread2.sig_get_img_1.connect(self.on_get_img_1)

            # 设置硬件地址
            self.xdma_driver.write_user_register(0x18, self.cam_addr_0)
            self.xdma_driver.write_user_register(0x1c, self.cam_addr_1)

            self.status_label.setText("硬件初始化完成")

        except Exception as e:
            self.status_label.setText(f"硬件初始化失败: {str(e)}")
            QMessageBox.warning(self, "错误", f"硬件初始化失败: {str(e)}")

    @Slot()
    def on_start_camera_clicked(self):
        """启动摄像头按钮点击事件"""
        if not self.start_flag:
            self.start_camera()
        else:
            self.stop_camera()

    def start_camera(self):
        """启动摄像头"""
        if not self.xdma_driver:
            QMessageBox.warning(self, "错误", "硬件未初始化")
            return

        try:
            # 启动摄像头采集
            self.xdma_driver.start_camera()

            # 启动所有线程（按需启动模式，匹配C++版本）
            if not self.event_thread.isRunning():
                self.event_thread.start()
            if not self.event_thread1.isRunning():
                self.event_thread1.start()
            if not self.event_thread2.isRunning():
                self.event_thread2.start()
            if not self.start_calc_thread.isRunning():
                self.start_calc_thread.start()

            self.start_flag = True
            self.start_camera_btn.setText("停止摄像头")
            self.status_label.setText("摄像头运行中...")

        except Exception as e:
            QMessageBox.warning(self, "错误", f"启动摄像头失败: {str(e)}")

    def stop_camera(self):
        """停止摄像头"""
        try:
            # 停止摄像头采集
            if self.xdma_driver:
                self.xdma_driver.stop_camera()

            # 停止所有线程
            if self.event_thread and self.event_thread.isRunning():
                self.event_thread.running = False
                self.event_thread.wait()

            if self.event_thread1 and self.event_thread1.isRunning():
                self.event_thread1.running = False
                self.event_thread1.wait()

            if self.event_thread2 and self.event_thread2.isRunning():
                self.event_thread2.running = False
                self.event_thread2.wait()

            if self.load_img_thread and self.load_img_thread.isRunning():
                self.load_img_thread.running = False
                self.load_img_thread.wait()

            if self.start_calc_thread and self.start_calc_thread.isRunning():
                self.start_calc_thread.running = False
                self.start_calc_thread.wait()

            self.start_flag = False
            self.start_camera_btn.setText("启动摄像头")
            self.status_label.setText("摄像头已停止")

        except Exception as e:
            QMessageBox.warning(self, "错误", f"停止摄像头失败: {str(e)}")

    @Slot()
    def on_load_instruction_clicked(self):
        """加载指令按钮点击事件"""
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(
            self, "选择指令文件", "", "Text files (*.txt);;All files (*.*)"
        )

        if file_path:
            try:
                # 这里应该加载指令文件内容
                # 暂时只显示加载成功
                self.status_label.setText(f"指令文件已加载: {os.path.basename(file_path)}")
            except Exception as e:
                QMessageBox.warning(self, "错误", f"加载指令文件失败: {str(e)}")

    @Slot()
    def on_load_weight_clicked(self):
        """加载权重按钮点击事件"""
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(
            self, "选择权重文件", "", "Binary files (*.bin);;All files (*.*)"
        )

        if file_path:
            try:
                # 这里应该加载权重文件内容
                # 暂时只显示加载成功
                self.status_label.setText(f"权重文件已加载: {os.path.basename(file_path)}")
            except Exception as e:
                QMessageBox.warning(self, "错误", f"加载权重文件失败: {str(e)}")

    @Slot()
    def on_get_img(self):
        """获取图像信号处理 - 对应C++的get_img()"""
        try:
            if not self.xdma_driver:
                return

            # 写入寄存器触发摄像头0
            self.xdma_driver.write_user_register(0x20, 1)

            # 从摄像头0读取图像数据
            image_data = self.xdma_driver.read_data(self.cam_addr_0, self.img_height * self.img_width * 3)

            if len(image_data) == self.img_height * self.img_width * 3:
                # 转换图像数据为OpenCV格式
                image_array = np.frombuffer(image_data, dtype=np.uint8)
                image_array = image_array.reshape((self.img_height, self.img_width, 3))

                # BGR转RGB（如果需要）
                image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)

                # 设置到SendImage对象中
                self.send_image.capture = image_array

                # 如果load_img_thread没有运行，则启动它
                if not self.load_img_thread.isRunning():
                    self.load_img_thread.start()

                # 重新启动event_thread1（匹配C++版本的逻辑）
                if not self.event_thread1.isRunning():
                    self.event_thread1.start()

        except Exception as e:
            print(f"获取图像失败: {e}")

    @Slot()
    def on_get_img_1(self):
        """获取图像1信号处理 - 对应C++的get_img_1()"""
        try:
            if not self.xdma_driver:
                return

            # 写入寄存器触发摄像头1
            self.xdma_driver.write_user_register(0x24, 1)

            # 从摄像头1读取图像数据
            image_data = self.xdma_driver.read_data(self.cam_addr_1, self.img_height * self.img_width * 3)

            if len(image_data) == self.img_height * self.img_width * 3:
                # 转换图像数据为OpenCV格式
                image_array = np.frombuffer(image_data, dtype=np.uint8)
                image_array = image_array.reshape((self.img_height, self.img_width, 3))

                # BGR转RGB（如果需要）
                image_array = cv2.cvtColor(image_array, cv2.COLOR_BGR2RGB)

                # 设置到SendImage对象中
                self.send_image.capture = image_array

                # 重新启动event_thread2（匹配C++版本的逻辑）
                if not self.event_thread2.isRunning():
                    self.event_thread2.start()

        except Exception as e:
            print(f"获取图像1失败: {e}")

    @Slot()
    def on_nms(self):
        """NMS处理信号 - 对应C++的nms()"""
        try:
            if not self.xdma_driver:
                return

            # 写入寄存器触发NMS
            self.xdma_driver.write_user_register(0x10, 1)

            # 获取处理后的图像用于显示
            if self.send_image.img3 is not None:
                self.current_image = self.send_image.img3.copy()
                self.update_image_display()

            # 重新启动event_thread（匹配C++版本的逻辑）
            if not self.event_thread.isRunning():
                self.event_thread.start()

        except Exception as e:
            print(f"NMS处理失败: {e}")


    def update_image_display(self):
        """更新图像显示"""
        if self.current_image is not None:
            # 转换图像格式用于显示
            rgb_image = cv2.cvtColor(self.current_image, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w

            # 创建QImage
            qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format.Format_RGB888)

            # 创建QPixmap并缩放到标签大小
            pixmap = QPixmap.fromImage(qt_image)
            scaled_pixmap = pixmap.scaled(
                self.image_label.size(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )

            # 显示图像
            self.image_label.setPixmap(scaled_pixmap)

    def closeEvent(self, event):
        """窗口关闭事件"""
        # 停止摄像头和线程
        self.stop_camera()

        # 停止所有线程
        threads = [self.event_thread, self.event_thread1, self.event_thread2,
                  self.load_img_thread, self.start_calc_thread]

        for thread in threads:
            if thread and thread.isRunning():
                thread.running = False
                thread.wait(1000)  # 等待1秒

        # 释放硬件资源
        if self.xdma_driver:
            self.xdma_driver.close()

        event.accept()
